//===--- ArgumentExplosionTransform.cpp -----------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2018 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
///
/// \file
///
/// This file contains an implementation of the partial dead argument
/// elimination optimization. We do this to attempt to remove non-trivial
/// arguments of callees to eliminate lifetime constraints of a large argument
/// on values in the caller.
///
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "fso-argument-explosion-transform"

#include "polarphp/pil/optimizer/internal/funcsignaturetransforms/FunctionSignatureOpts.h"
#include "llvm/Support/CommandLine.h"

using namespace polar;

static llvm::cl::opt<bool> FSODisableArgExplosion(
   "sil-fso-disable-arg-explosion",
   llvm::cl::desc("Do not perform argument explosion during FSO. Intended "
                  "only for testing purposes"));

//===----------------------------------------------------------------------===//
//                                  Utility
//===----------------------------------------------------------------------===//

/// Whether the known-to-date upper bound on the live leaf count is high enough
/// so that argument explosion is possible.
static bool
mayExplodeGivenLiveLeafCountUpperBound(unsigned knownLiveLeafCountUpperBound) {
   return knownLiveLeafCountUpperBound > 0;
}

static unsigned maxExplosionSizeWhenSpecializationWillIntroduceThunk(
   bool willSpecializationIntroduceThunk) {
   // 3 is the heuristic max explosion size for a single argument when the
   // specializing the function will introduce a thunk.  If specializing the
   // function may not introduce a thunk, then we rely on the maximum size
   // imposed by shouldExpand.
   return willSpecializationIntroduceThunk ? 3 : UINT_MAX;
}

static bool shouldExplode(unsigned knownLiveLeafCountUpperBound,
                          bool hasKnownDeadLeaves,
                          bool hasKnownDeadNontrivialLeaves,
                          bool willSpecializationIntroduceThunk) {
   unsigned maxExplosionSize =
      maxExplosionSizeWhenSpecializationWillIntroduceThunk(
         /*willSpecializationIntroduceThunk=*/
         willSpecializationIntroduceThunk);
   bool isLiveLeafCountInExplodableRange =
      mayExplodeGivenLiveLeafCountUpperBound(knownLiveLeafCountUpperBound) &&
      (knownLiveLeafCountUpperBound <= maxExplosionSize);
   bool hasKnownDeadRelevantLeaves = willSpecializationIntroduceThunk
                                     ? hasKnownDeadNontrivialLeaves
                                     : hasKnownDeadLeaves;
   return isLiveLeafCountInExplodableRange && hasKnownDeadRelevantLeaves;
}

/// Return true if it's both legal and a good idea to explode this argument.
///
/// Our main interest here is to expose more opportunities for ARC. This means
/// that we are not interested in exploding (and partially DCEing) structs in
/// the following cases:
///
/// 1. Completely dead arguments. This is handled by dead argument elimination.
///
/// 2. Structs with many live leaf nodes. Our heuristic is to explode if there
///    are only 1-3 live leaf nodes for specializations and 1-6 live leaf nodes
///    (in fact, the number specified in shouldExpand). Otherwise again we run
///    into register pressure/spilling issues.
///    TODO: Improve the 1-3 heuristic by having FSO consider the total
///          resultant argument count.  Currently, there is no consideration of
///          that, meaning we could end up with argument exploding even in the
///          case of long argument lists where it isn't beneficial.
///
/// Perform argument exploding if one of the following sets of conditions hold:
///
/// 1. a. The live leaf count is less than or equal to 3.
///    b. There is a dead non-trivial leaf.
/// 2. a. The live leaf count is less than or equal to 6.
///    b. There is a dead trivial leaf.
///    c. Specializing the function will not result in a thunk.
static bool
shouldExplode(FunctionSignatureTransformDescriptor &transformDesc,
              ArgumentDescriptor &argDesc,
              ConsumedArgToEpilogueReleaseMatcher &epilogueReleaseMatcher) {
   // The method is structured as follows:
   //
   // First, do some basic checks and exit early.
   // Then in three steps of increasing complexity, calculate data which could
   // permit the heuristic to decide to explode the argument.  These steps
   // provide information of increasing expense and fidelity.  Checking whether
   // the heuristic allows explosion after each step unnecessary work to be
   // avoided.
   //
   // In a bit more detail:
   //
   // 1) Do some basic checks and exit early, returning false.
   //    - that we can optimize the argument at all
   //    - that the argument has more than a single leaf node
   //    - that the module permits the type to be expanded
   // 2) Gather some basic leaf counts.
   //    - calculate the unmodified (unmodified that is by the results of the
   //      owned-to-guaranteed transformation) live leaf count
   //    - calculate the total list of leaf types to obtain the total leaf count
   // 3) Check whether the heuristic allows the argument to be exploded using
   //    only potentially-trivial leaf counts.  At this point it is certainly not
   //    known that there are dead non-trivial leaves, so exiting early here
   //    is only possible if specializing the function will not result in a
   //    thunk.
   // 4) Gather the counts of non-trivial leaves.
   //    - calculate the count of total non-trivial leaves by filtering the total
   //      list of leaf types from step 2) according to whether leaf is trivial
   //    - calculate an upper bound (upper bound because it doesn't consider the
   //      results of the owned-to-guaranteed transformation) on the count of
   //      live non-trivial leaves
   // 5) Check whether the heuristic allows the argument to be exploded using the
   //    upper bound on live non-trivial leaves.
   // 6) Dial in the upper bounds calculated in steps 2) and 4) by compensating
   //    for the effects of the owned-to-guaranteed transformation.
   // 7) Check whether the heuristic allows the argument to be exploded using the
   //    actual count of live leaves, both trivial and non-trivial.

   // No passes can optimize this argument, so just bail.
   if (!argDesc.canOptimizeLiveArg()) {
      LLVM_DEBUG(llvm::dbgs()
                    << "The argument is of a type that cannot be exploded.");
      return false;
   }

   // If the argument is a singleton, it will not be exploded.
   //
   // Explosion makes sense only if some but not all of the leaves are live.
   //
   // Note that ProjectionTree::isSingleton returns true for enums since they are
   // sums and not products and so only have a single top-level node.
   if (argDesc.ProjTree.isSingleton()) {
      LLVM_DEBUG(llvm::dbgs() << "The argument's type is a singleton.");
      return false;
   }

   auto *argument = argDesc.Arg;
   auto &module = argument->getModule();
   auto type = argument->getType().getObjectType();

   // If the global type expansion heuristic does not allow the type to be
   // expanded, it will not be exploded.
   if (!shouldExpand(module, type)) {
      LLVM_DEBUG(llvm::dbgs()
                    << "The argument is of a type which should not be expanded.");
      return false;
   }

   bool willSpecializationIntroduceThunk =
      transformDesc.willSpecializationIntroduceThunk();

   unsigned const liveLeafCountUpperBound = argDesc.ProjTree.getLiveLeafCount();

   // If we know already that we may not explode given the upper bound we have
   // established on the live leaf count, exit early.
   //
   // If the argument is completely dead, it will not be exploded.
   //
   // Explosion makes sense only if some but not all of the leaves are live.  The
   // dead argument transformation will try to eliminate the argument.
   if (!mayExplodeGivenLiveLeafCountUpperBound(liveLeafCountUpperBound)) {
      LLVM_DEBUG(llvm::dbgs() << "The argument has no live leaves.");
      return false;
   }

   // To determine whether some but not all of the leaves are used, the total
   // leaf count must be retrieved.
   llvm::SmallVector<PILType, 32> allLeaves;
   argDesc.ProjTree.getAllLeafTypes(allLeaves);
   unsigned const leafCount = allLeaves.size();

   assert(
      liveLeafCountUpperBound <= leafCount &&
      "There should be no more *live* leaves than there are *total* leaves.");

   if (shouldExplode(
      /*knownLifeLeafCount=*/liveLeafCountUpperBound,
      /*hasKnownDeadLeaves=*/liveLeafCountUpperBound < leafCount,
      /*hasKnownDeadNontrivialLeaves=*/false,
      /*willSpecializationIntroduceThunk=*/
                             willSpecializationIntroduceThunk)) {
      LLVM_DEBUG(
         llvm::dbgs()
            << "Without considering the liveness of non-trivial leaves, it has "
               "already been determined that there are already fewer ("
            << liveLeafCountUpperBound
            << ") live leaves of the relevant sort (trivial) than total leaves ("
            << leafCount << ") and no more total live leaves ("
            << liveLeafCountUpperBound << ") than the heuristic permits ("
            << maxExplosionSizeWhenSpecializationWillIntroduceThunk(
               /*willSpecializationIntroduceThunk=*/
               willSpecializationIntroduceThunk)
            << ").  Exploding.");
      return true;
   }

   auto *function = argument->getFunction();
   unsigned const nontrivialLeafCount = llvm::count_if(
      allLeaves, [&](PILType type) { return !type.isTrivial(*function); });

   llvm::SmallVector<const ProjectionTreeNode *, 32> liveLeaves;
   argDesc.ProjTree.getLiveLeafNodes(liveLeaves);
   // NOTE: The value obtained here is an upper bound because the
   //       owned-to-guaranteed transformation may eliminate some live
   //       non-trivial leaves, leaving the count lower.
   unsigned const liveNontrivialLeafCountUpperBound =
      llvm::count_if(liveLeaves, [&](const ProjectionTreeNode *leaf) {
         return !leaf->getType().isTrivial(*function);
      });

   assert(liveNontrivialLeafCountUpperBound <= nontrivialLeafCount &&
          "There should be no more *live* non-trivial leaves than there are "
          "*total* non-trivial leaves.");
   assert(nontrivialLeafCount <= leafCount &&
          "There should be no more *non-trivial* leaves than there are *total* "
          "leaves.");

   // If it is known without taking the owned-to-guaranteed transformation into
   // account both that exploding will reduce ARC traffic (because an upper bound
   // for the number of live non-trivial leaves is less than the non-trivial
   // leaf count) and also that the explosion will fit within the heuristic upper
   // bound (because an upper bound for the total live leaf count falls within
   // the limit imposed by the heuristic), then explode now.
   bool shouldExplodeGivenUpperBounds = shouldExplode(
      /*knownLiveLeafCount=*/liveLeafCountUpperBound,
      /*hasKnownDeadLeaves=*/liveLeafCountUpperBound < leafCount,
      /*hasKnownDeadNontrivialLeaves=*/liveNontrivialLeafCountUpperBound <
                                       nontrivialLeafCount,
      /*willSpecializationIntroduceThunk=*/willSpecializationIntroduceThunk);
   if (shouldExplodeGivenUpperBounds) {
      LLVM_DEBUG(
         llvm::dbgs()
            << "Without considering the expected results of the "
               "owned-to-guaranteed transformation, there are already fewer ("
            << liveNontrivialLeafCountUpperBound
            << ") live non-trivial leaves than total leaves ("
            << nontrivialLeafCount << ") and no more total live leaves ("
            << liveLeafCountUpperBound << ") than the heuristic permits ("
            << maxExplosionSizeWhenSpecializationWillIntroduceThunk(
               /*willSpecializationIntroduceThunk=*/
               willSpecializationIntroduceThunk)
            << ").  Exploding.");
      return true;
   }

   unsigned liveLeafCount = liveLeafCountUpperBound;
   unsigned liveNontrivialLeafCount = liveNontrivialLeafCountUpperBound;

   // The upper bounds that have been established for the live leaf counts are
   // too high to permit us to explode.  That could be because it hasn't been
   // established that any leaves are dead or alternatively that it hasn't been
   // established that there are fewer total live leaves than the limit imposed
   // by the heuristic.  In either case, if some live leaves are eliminated, the
   // number of live leaves may decrease such that exploding will be possible.
   // The results of the owned-to-guaranteed transformation are predicated.  If
   // it is predicted that a leaf will be dead after the owned-to-guaranteed
   // transformation, then the leaf count is decreased.
   //
   // The owned-to-guaranteed will only be applied to the argumehnt if its
   // convention is Direct_Owned.  Additionally, it only applies to non-trivial
   // leaves, which it may kill, so if it is already known that there are no live
   // non-trivial leaves, owned-to-guaranteed will not eliminate anything.
   if (argDesc.hasConvention(PILArgumentConvention::Direct_Owned) &&
       liveNontrivialLeafCountUpperBound > 0) {
      if (auto maybeReleases =
         epilogueReleaseMatcher.getPartiallyPostDomReleaseSet(argument)) {
         auto releases = maybeReleases.getValue();
         llvm::SmallPtrSet<PILInstruction *, 8> users;
         users.insert(std::begin(releases), std::end(releases));

         for (auto *leaf : liveLeaves) {
            if (llvm::all_of(leaf->getNonProjUsers(), [&](Operand *operand) {
               return users.count(operand->getUser());
            })) {
               // Every non-projection user of the leaf is an epilogue release.  The
               // owned-to-guaranteed transformation will eliminate this usage.  With
               // the expectation of that usage being eliminated, stop considering
               // this leaf to be live for the purposes of deciding whether the
               // argument should be exploded.
               --liveLeafCount;
               --liveNontrivialLeafCount;
            }
         }
      }
   }

   return shouldExplode(
      /*knownLifeLeafCount=*/liveLeafCount,
      /*hasKnownDeadLeaves=*/liveLeafCount < leafCount,
      /*hasKnownDeadNontrivialLeaves=*/liveNontrivialLeafCount <
                                       nontrivialLeafCount,
      /*willSpecializationIntroduceThunk=*/willSpecializationIntroduceThunk);
}

//===----------------------------------------------------------------------===//
//                          Top Level Implementation
//===----------------------------------------------------------------------===//

bool FunctionSignatureTransform::ArgumentExplosionAnalyzeParameters() {
   // If we are not supposed to perform argument explosion, bail.
   if (FSODisableArgExplosion)
      return false;

   PILFunction *F = TransformDescriptor.OriginalFunction;
   // Did we decide we should optimize any parameter?
   bool SignatureOptimize = false;
   auto Args = F->begin()->getPILFunctionArguments();
   ConsumedArgToEpilogueReleaseMatcher ArgToReturnReleaseMap(
      RCIA->get(F), F, {PILArgumentConvention::Direct_Owned});

   // Analyze the argument information.
   for (unsigned i : indices(Args)) {
      ArgumentDescriptor &A = TransformDescriptor.ArgumentDescList[i];
      // If the argument is dead, there is no point in trying to explode it. The
      // dead argument pass will get it.
      if (A.IsEntirelyDead) {
         continue;
      }

      // Do not optimize argument.
      if (!A.canOptimizeLiveArg()) {
         continue;
      }

      // Explosion of generic parameters is not supported yet.
      if (A.Arg->getType().hasArchetype())
         continue;

      A.ProjTree.computeUsesAndLiveness(A.Arg);
      A.Explode = shouldExplode(TransformDescriptor, A, ArgToReturnReleaseMap);

      // Modified self argument.
      if (A.Explode && Args[i]->isSelf()) {
         TransformDescriptor.shouldModifySelfArgument = true;
      }

      SignatureOptimize |= A.Explode;
   }
   return SignatureOptimize;
}

void FunctionSignatureTransform::ArgumentExplosionFinalizeOptimizedFunction() {
   PILFunction *NewF = TransformDescriptor.OptimizedFunction.get();
   PILBasicBlock *BB = &*NewF->begin();
   PILBuilder Builder(BB->begin());
   Builder.setCurrentDebugScope(BB->getParent()->getDebugScope());
   unsigned TotalArgIndex = 0;
   for (ArgumentDescriptor &AD : TransformDescriptor.ArgumentDescList) {
      // If this argument descriptor was dead and we removed it, just skip it. Do
      // not increment the argument index.
      if (AD.WasErased) {
         continue;
      }

      // Simply continue if do not explode.
      if (!AD.Explode) {
         TransformDescriptor.AIM[TotalArgIndex] = AD.Index;
         ++TotalArgIndex;
         continue;
      }

      assert(!AD.IsEntirelyDead &&
             "Should never see completely dead values here");

      // OK, we need to explode this argument.
      unsigned ArgOffset = ++TotalArgIndex;
      unsigned OldArgIndex = ArgOffset - 1;
      llvm::SmallVector<PILValue, 8> LeafValues;

      // We do this in the same order as leaf types since ProjTree expects that
      // the order of leaf values matches the order of leaf types.
      llvm::SmallVector<const ProjectionTreeNode *, 8> LeafNodes;
      AD.ProjTree.getLiveLeafNodes(LeafNodes);

      for (auto *Node : LeafNodes) {
         auto OwnershipKind = *AD.getTransformedOwnershipKind(Node->getType());
         LeafValues.push_back(
            BB->insertFunctionArgument(ArgOffset, Node->getType(), OwnershipKind,
                                       BB->getArgument(OldArgIndex)->getDecl()));
         TransformDescriptor.AIM[TotalArgIndex - 1] = AD.Index;
         ++ArgOffset;
         ++TotalArgIndex;
      }

      // Then go through the projection tree constructing aggregates and replacing
      // uses.
      AD.ProjTree.replaceValueUsesWithLeafUses(
         Builder, BB->getParent()->getLocation(), LeafValues);

      // We ignored debugvalue uses when we constructed the new arguments, in
      // order to preserve as much information as possible, we construct a new
      // value for OrigArg from the leaf values and use that in place of the
      // OrigArg.
      PILValue NewOrigArgValue = AD.ProjTree.computeExplodedArgumentValue(
         Builder, BB->getParent()->getLocation(), LeafValues);

      // Replace all uses of the original arg with the new value.
      PILArgument *OrigArg = BB->getArgument(OldArgIndex);
      OrigArg->replaceAllUsesWith(NewOrigArgValue);

      // Now erase the old argument since it does not have any uses. We also
      // decrement ArgOffset since we have one less argument now.
      BB->eraseArgument(OldArgIndex);
      --TotalArgIndex;
   }
}
